from calc_manager.plugins_inner.bdf_reader_inner import BDFReader
from calc_manager.plugins_inner.op2_reader_inner import OP2Reader
from calc_manager.plugins_inner.mat_csv_inner import MatCSVReader
from calc_manager.plugins_inner.bdf_transformer_inner import BDFTransformer
from pydantic import BaseModel


class Parameter(BaseModel):
    sample_num: int 
    sample_float: float

class KeyResult(BaseModel):
    y: float
    rrr: str

refine_reference_eids = [
    list(range(10000001, 15000000)),
    list(range(15000001, 20000000)),
    list(range(20000001, 29999999)),
    list(range(30000001, 36999999)),
    list(range(37000001, 39999999)),
    list(range(40000001, 44999999)),
    list(range(50000001, 54999999)),
    list(range(60000001, 64999999)),
    list(range(70000001, 79999999)),
    list(range(80000001, 89999999)),
    list(range(90000001, 94999999)),
    list(range(91000001, 91499999)),
    list(range(91500001, 91999999)),
    list(range(92000001, 92999999)),
    list(range(93000001, 93699999)),
    list(range(93700001, 93999999)),
    list(range(94000001, 94499999)),
    list(range(95000001, 95499999)),
    list(range(96000001, 96499999)),
    list(range(97000001, 97999999)),
    list(range(98000001, 98999999)),
    list(range(99000001, 99499999)),
]

def check_eid(eids):
    set_refine_eids = set([item for subset in refine_reference_eids for item in subset])
    # is_feine: 1代表精细网格，0代表自然网格
    # result = {
    #     "eid": [],
    #     "is_refine":[]
    # }
    # for eid in eids:
    #     result["eid"].append(eid)
    #     if eid in set_refine_eids:
    #         result["is_refine"].append(1)
    #     else:
    #         result["is_refine"].append(0)
    rough_eids = [item for item in eids if item not in set_refine_eids]
    refine_eids = [item for item in eids if item in set_refine_eids]
    # 前者代表自然网格id,后者代表精细网格id
    return rough_eids, refine_eids



def refine_main(param: Parameter, bdf_reader: BDFReader, op2_reader: OP2Reader, mat_csv_reader: MatCSVReader, bdf_transformer: BDFTransformer) :
    print("refine_main")
    eids = bdf_reader.get_element_ids()
    rough_eids, refine_eids = check_eid(eids)
    #print(f"rough_eids: {rough_eids}, refine_eids: {refine_eids}")

    results = {}
    subcase_id = op2_reader.get_subcase_id()
    result_components = ["stress.ctria3_stress", "stress.cquad4_stress"]
    for comp in result_components:
        print(f"Getting results for component {comp} and subcase {subcase_id}")
        stress = op2_reader.get_result_from_comp_and_subcase(comp, subcase_id)
        print(f"Got stress results with shape: {stress.data.shape}")
        element_node = stress.element_node
        von_mises = stress.von_mises()
        for index, item in enumerate(element_node):
            if index % 2 == 0:
                results[item[0]] = von_mises[0][index : index + 2].tolist()

    # 删除自然网格的结果(无需关注)
    for eid in rough_eids:
        results.pop(eid, None)

    """
    1. 失效单元eid
    2. 失效单元eid对应的cid,对应cname
    3. 根据cname获取到g70name,再分组
    4. 第二张图: 根据g70name分组, 补充当前失效cnames包含的所有eids
    6. 第三张图: 根据g70name分组, 补充当前失效g70下所有的cnames的所有eids
    """
    ## TODO 理论上需要遍历所有的bdf文件获取所有数据,目前暂时指定文件名
    #refine_bdf_path = r"D:\SF-CAE\model\S20.bdf"  # 需要替换为真实的精细网格bdf文件
    #bdf_transformer = BDFTransformer(refine_bdf_path)
    comps_info = bdf_transformer.get_comps_info()

    # 获取所有eid对应的安全系数以及失效单元eid
    # danger_infos包含g70_name和对应的cids关系
    # 组织层级 gname-->cname-->eid
    danger_infos = {}
    for refine_eid in refine_eids:
        element = bdf_reader.get_element(refine_eid)
        mat_id = element.pid_ref.material_ids[0]
        # 根据mat_id获取材料许用应力数据
        material_ultimate_stress = mat_csv_reader.get_mat_from_csv(mat_id)
        element_stress = results[refine_eid]
        element_stress.append(material_ultimate_stress / 3)
        coffe = material_ultimate_stress / max(element_stress) - 1
        if coffe < 0:
            # 获取cid
            cid = bdf_transformer.get_comp_by_eid(refine_eid)
            for comp_info in comps_info:
                if cid == comp_info["cid"]:
                    cname = comp_info["cname"]
                    # 上航提供的名称规则
                    g70_name = cname.split("_")[0]
                    break
            if g70_name not in danger_infos:
                danger_infos[g70_name] = [cid]
            elif cid not in danger_infos[g70_name]:
                danger_infos[g70_name].append(cid)
        results[refine_eid].append(coffe)

    print(f"danger_infos: {danger_infos}")

    """
    构造最终裕度图需要的数据
    1. output长度代表最终报表的行数
    2. gname代表表格第一列的G70编号结果
    3. 基于eids和coffe三维负责渲染数据
    4. num代表表格第三列的失效单元数
    """
    margin_output = []
    for k, v in danger_infos.items():
        eids = []
        for cid in v:
            eids.extend(bdf_transformer.get_comp_eids_ref()[cid])
        coffe = [results[eid][-1] for eid in eids]
        num = sum(1 for i in coffe if i < 0)
        margin_output.append(
            {
                "gname": k,
                "eids": eids,
                "coffe": coffe,
                "num": num,
            }
        )

    """
    第三张云图,只需要g70_name的数据
    1. gname代表表格第一列的G70编号结果
    2. 基于eids和stress数据绘制第三张应力云图
    """
    # TODO 这个数据可以和第二张图的数据合并
    stress_output = []
    for k, v in danger_infos.items():
        eids = bdf_transformer.get_comps_by_gname(k)[1]
        stress = [max(results[eid]) for eid in eids]
        stress_output.append(
            {
                "gname": k,
                "eids": eids,
                "stress": stress
            }
        )

    #### 如果通过算子算法平台计算,则需要返回以下数据
    """
    1. rough_eids, refine_eids
    2. margin_output
    3. stress_output
    """

    res =  KeyResult(y=3.14, rrr="rrr")
    return res
